from manim import *

class ArrayElement(VGroup):
    """表示数组中的单个元素，包含矩形和标签"""
    def __init__(self, value, max_value, **kwargs):
        super().__init__(**kwargs)
        # 创建矩形
        height = value / max_value * 3
        self.rect = Rectangle(height=height, width=0.6, fill_opacity=0.8, fill_color=BLUE, stroke_color=WHITE)
        # 创建标签
        self.label = Text(str(value), font="SimHei").scale(0.6)
        self.label.next_to(self.rect, DOWN, buff=0.1)
        # 将矩形和标签添加到组中
        self.add(self.rect, self.label)
        self.value = value

class MergeSortVisualization(Scene):
    def construct(self):
        # 创建中央标题并显示1秒后消失
        title = Text("归并排序", font="SimHei").scale(1.5)
        title.move_to(ORIGIN)  # 放置在屏幕中央
        
        self.play(FadeIn(title))
        self.wait(1)
        self.play(FadeOut(title))
        
        # 创建待排序的数组
        array = [7, 3, 1, 9, 5, 2, 8, 6]
        
        # 创建数组元素
        elements = self.create_array_elements(array)
        elements_group = VGroup(*elements).arrange(RIGHT, buff=0.2)
        elements_group.to_edge(UP, buff=1.5)  # 放在上方，留出足够空间给中央标题
        
        # 添加数组到场景
        self.play(FadeIn(elements_group))
        self.wait(0.5)
        
        # 开始归并排序可视化
        sorted_array, sorted_elements = self.merge_sort_animate(array, elements)
        
        # 突出显示排序完成
        self.play(
            *[element.rect.animate.set_color(GREEN) for element in sorted_elements],
            run_time=1.5,
            rate_func=smooth
        )
        self.wait(1)
        
        # 显示"排序完成"文本
        completion_text = Text("排序完成", font="SimHei", color=GREEN).scale(1.2)
        completion_text.to_edge(DOWN)
        self.play(FadeIn(completion_text))
        self.wait(2)
    
    def create_array_elements(self, array):
        """创建数组元素"""
        max_value = max(array)
        return [ArrayElement(value, max_value) for value in array]
    
    def merge_sort_animate(self, array, elements, level=0, offset=0):
        """递归地可视化归并排序算法"""
        n = len(array)
        
        # 基本情况：只有一个元素时，无需排序
        if n <= 1:
            return array, elements

        # 计算中间位置
        mid = n // 2
        
        # 创建左右子数组的副本
        left_array = array[:mid]
        right_array = array[mid:]
        
        left_elements = elements[:mid]
        right_elements = elements[mid:]
        
        # 可视化分割过程
        if level == 0:
            # 分割阶段文字放在屏幕中央
            split_text = Text("分割阶段", font="SimHei", color=YELLOW).scale(1.2)
            split_text.move_to(ORIGIN).to_edge(UP, buff=0.3)  # 放在上方中央
            self.play(FadeIn(split_text))
            self.wait(0.5)
        
        # 移动元素到新的位置来表示分割
        level_shift = 0.4 * (level + 1) + 0.4
        
        left_target = VGroup(*left_elements).arrange(RIGHT, buff=0.2)
        right_target = VGroup(*right_elements).arrange(RIGHT, buff=0.2)
        
        # 计算水平偏移
        left_offset = offset - (n / 4) * 0.8
        right_offset = offset + (n / 4) * 0.8
        
        left_target.shift(DOWN * level_shift + RIGHT * left_offset)
        right_target.shift(DOWN * level_shift + RIGHT * right_offset)
        
        # 动画显示分割
        if n > 2:
            self.play(
                *[element.animate.move_to(target) for element, target in zip(left_elements, left_target)],
                *[element.animate.move_to(target) for element, target in zip(right_elements, right_target)],
                run_time=1.5,
                rate_func=smooth
            )
            self.wait(0.5)
        
        # 递归排序左右子数组
        sorted_left, sorted_left_elements = self.merge_sort_animate(
            left_array, left_elements, level + 1, left_offset
        )
        
        sorted_right, sorted_right_elements = self.merge_sort_animate(
            right_array, right_elements, level + 1, right_offset
        )
        
        # 如果需要，可以在递归返回后添加位置调整的动画
        if level == 0:
            # 分割阶段结束，可以添加数组元素的过渡动画
            self.play(
                *[elem.animate for elem in sorted_left_elements + sorted_right_elements],
                run_time=0.8,
                rate_func=smooth
            )
        
        # 可视化合并过程
        if level == 0:
            # 分割阶段文字淡出，合并阶段文字淡入
            split_text_out = split_text
            merge_text = Text("合并阶段", font="SimHei", color=GREEN).scale(1.2)
            merge_text.move_to(split_text.get_center())  # 放在分割阶段文字的位置上
            self.play(FadeOut(split_text_out), FadeIn(merge_text))
            self.wait(0.5)
        
        # 合并左右子数组
        merged, merged_elements = self.merge_animate(
            sorted_left, sorted_left_elements,
            sorted_right, sorted_right_elements,
            level, offset
        )
        
        return merged, merged_elements
    
    def merge_animate(self, left, left_elements, right, right_elements, level, offset):
        """可视化两个已排序数组的合并过程"""
        merged = []
        merged_elements = []
        
        i = j = 0
        
        # 计算合并后的位置
        target_y = -0.8 * level if level > 0 else 0
        
        while i < len(left) and j < len(right):
            # 高亮比较的两个元素
            if i < len(left) and j < len(right):
                self.play(
                    left_elements[i].rect.animate.set_color(YELLOW),
                    right_elements[j].rect.animate.set_color(YELLOW),
                    run_time=0.5,
                    rate_func=smooth
                )
                self.wait(0.3)
            
            if left[i] <= right[j]:
                merged.append(left[i])
                merged_elements.append(left_elements[i])
                
                # 移动到合并数组的位置
                target_x = offset + (len(merged) - 1 - (len(left) + len(right)) / 2) * 0.8
                self.play(
                    left_elements[i].animate.move_to([target_x, target_y, 0]),
                    run_time=0.8,
                    rate_func=smooth
                )
                i += 1
            else:
                merged.append(right[j])
                merged_elements.append(right_elements[j])
                
                # 移动到合并数组的位置
                target_x = offset + (len(merged) - 1 - (len(left) + len(right)) / 2) * 0.8
                self.play(
                    right_elements[j].animate.move_to([target_x, target_y, 0]),
                    run_time=0.8,
                    rate_func=smooth
                )
                j += 1
        
        # 处理剩余元素
        while i < len(left):
            merged.append(left[i])
            merged_elements.append(left_elements[i])
            
            target_x = offset + (len(merged) - 1 - (len(left) + len(right)) / 2) * 0.8
            self.play(
                left_elements[i].animate.move_to([target_x, target_y, 0]),
                run_time=0.8,
                rate_func=smooth
            )
            i += 1
        
        while j < len(right):
            merged.append(right[j])
            merged_elements.append(right_elements[j])
            
            target_x = offset + (len(merged) - 1 - (len(left) + len(right)) / 2) * 0.8
            self.play(
                right_elements[j].animate.move_to([target_x, target_y, 0]),
                run_time=0.8,
                rate_func=smooth
            )
            j += 1
        
        # 更新所有合并后元素的颜色
        for element in merged_elements:
            element.rect.set_color(BLUE)
        
        return merged, merged_elements 